package repeater

import (
	"context"
	"github.com/smallnest/rpcx/share"
	"github.com/xcy8712622040/gnetws/serverhandler"
	"sync"
)

// 元数据不是服务请求和服务响应的业务数据，而是一些辅助性的数据。
// 后端服务通过在响应中设置的元数据-会回写到这个里面，并且在每次请求中都会携带

var pool = sync.Pool{
	New: func() interface{} { return &Meta{mux: sync.RWMutex{}} },
}

type Meta struct {
	mux sync.RWMutex

	res map[string]string // rpcx 服务端返回的元数据
	req map[string]string // rpcx 请求元数据 (gateway中用于存储当前conn的元数据）
}

func (s *Meta) Set(key, val string) {
	s.mux.Lock()
	defer s.mux.Unlock()

	s.req[key] = val
	if s.res != nil {
		s.res[key] = val
	}
}

func (s *Meta) Del(key string) {
	s.mux.RLock()
	defer s.mux.RUnlock()

	delete(s.req, key)
	if s.res != nil {
		s.res[key] = ""
	}
}

func (s *Meta) Get(key string) string {
	s.mux.RLock()
	defer s.mux.RUnlock()

	return s.req[key]
}

func (s *Meta) Range(eval func(key, val string) bool) {
	s.mux.RLock()
	defer s.mux.RUnlock()

	for key, val := range s.req {
		if !eval(key, val) {
			return
		}
	}
}

func WaitReleaseForMeta(s *Meta) { s.res = nil; pool.Put(s) }

// PreMeta 从ctx中获取用户基本数据
func PreMeta(ctx context.Context) (s *Meta) {
	s = pool.Get().(*Meta)

	// rpcx service ctx 携带这个数据
	if val := ctx.Value(share.ResMetaDataKey); val != nil {
		s.res = val.(map[string]string)
	}

	// rpcx service ctx 与 gateway ctx 都会携带这个数据
	if val := ctx.Value(share.ReqMetaDataKey); val != nil {
		s.req = val.(map[string]string)
	} else {
		s.req = make(map[string]string)
		if gctx, ok := ctx.(*serverhandler.Context); ok {
			gctx.MetaData().SetInterface(share.ReqMetaDataKey, s.req)
		}
	}

	return
}
